package com.fanshuai.tree.trie;

import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;

import java.util.*;

/**
 * 前缀树。用于关键词匹配
 */
public class Trie implements ITrie {
    class Node {
        public boolean isEnd;
        public Map<Character, Node> subnodes;

        public Node() {
            isEnd = false;
            subnodes = new HashMap<>();
        }

        public void addSubNode(char ch, Node node) {
            subnodes.put(ch, node);
        }
        public void delSubNode(char ch) {
            subnodes.remove(ch);
        }

        public Node getSubNode(char ch) {
            return subnodes.get(ch);
        }
    }

    private Node root;
    private int size;

    public Trie() {
        root = new Node();
        size = 0;
    }

    @Override
    public void addWord(String word) {
        if (StringUtils.isBlank(word)) {
            return;
        }

        Node node = root;
        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            if (Character.isWhitespace(ch)) {
                continue;
            }

            Node sub = node.getSubNode(ch);
            if (null == sub) {
                sub = new Node();
                node.addSubNode(ch, sub);
            }
            node = sub;
            if (i == word.length() - 1) {
                sub.isEnd = true;
                size++;
            }
        }
    }

    @Override
    public void addWords(List<String> words) {
        if (CollectionUtils.isNotEmpty(words)) {
            words.forEach(this::addWord);
        }
    }

    @Override
    public List<String> getAllWords() {
        List<String> allWords = new ArrayList<>(size);
        StringBuilder b = new StringBuilder();

        fillWords(root, b, allWords);
        return allWords;
    }

    private void fillWords(Node node, StringBuilder b, List<String> words) {
        if (node.isEnd) {//完整关键词
            words.add(b.toString());
        }
        if (node.subnodes.size() == 0) {//叶节点
            return;
        }

        for (Character ch : node.subnodes.keySet()) {
            b.append(ch);
            fillWords(node.getSubNode(ch), b, words);

            int lastIndex = b.toString().length() - 1;
            b.deleteCharAt(lastIndex);
        }
    }

    @Override
    public List<WordToken> match(String text) {
        if (StringUtils.isBlank(text) || size == 0) {
            return new ArrayList<>();
        }

        List<WordToken> hitWords = new ArrayList<>();

        Node node = root;
        int i = 0;
        int k = i;
        while (i < text.length()) {
            char ch = text.charAt(i);
            if (Character.isWhitespace(ch)) {
                i++;
                continue;
            }

            node = node.getSubNode(ch);
            if (null == node) {
                i++;
                k = i;
                node = root;
            } else {
                if (node.isEnd) {
                    String word = text.substring(k, i + 1);
                    WordToken token = new WordToken(word, k, i);
                    hitWords.add(token);

                    i++;
                    k = i;
                    node = root;
                } else {
                    i++;
                }
            }
        }

        return hitWords;
    }

    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean isEmpty() {
        return size == 0;
    }
}
